package geppetto.cat.programs;


import geppetto.cat.alignments.Alignment;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.alignments.output.AlignerOutputLatex;
import geppetto.cat.common.StaticTools;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.AbstractModel;
import geppetto.cat.models.AgreementHMM;
import geppetto.cat.models.HMM;
import geppetto.cat.models.SubstochasticHMM;
import geppetto.phraseHMM.BijectiveHMM;
import geppetto.phraseHMM.RegularHMM;
import geppetto.phraseHMM.SymmetricHMM;

import java.io.IOException;
import java.io.PrintStream;
import java.util.Random;



public class PrettyPrintAlignments {

	
	String corpusDescription;

	int size, maxSentenceLen;
	String modelDir;
	String modelName;
	String baseDir;
	String decodingType;
	float posterior;	
	AbstractModel mhmm, mhmmb;
	BilingualCorpus forwardCorpus, backwardCorpus;

	double slack, epsilon;
	int maxStep, maxProjectionIterations;
	boolean projecTestTime;
	

	
	 public PrettyPrintAlignments (String[] args){
		corpusDescription = args[0];
		size = Integer.parseInt(args[1]);
		maxSentenceLen = Integer.parseInt(args[2]);
		modelDir = args[3];
		modelName = args[4];
		baseDir = args[5];
		projecTestTime = Boolean.parseBoolean(args[6]);
		decodingType = args[7];
		System.out.println("Corpus " + corpusDescription);
		System.out.println("Size " + size);
		System.out.println("Max Sentence Len " + maxSentenceLen);
		System.out.println("ModelDir " + modelDir);
		System.out.println("Model Name " + modelName);
		System.out.println("Output Dir " + baseDir);
		System.out.println("Project Decoding " + projecTestTime);
		System.out.println("Decoding type " + decodingType);
		System.out.println("---");
		if(modelName.equalsIgnoreCase("bijective") || modelName.equalsIgnoreCase("symmetric")){
			if (args.length < 11){
				System.out.println("Missing projection parameteres");
				System.out.println("Param: corpusFile Size MaxSentenceSize modelDir modelName ProjectAtDecode decodeType posteriorValue Epsilon Slack MaxProjectionStep MaxProjectionIteratiosn ");
			}
			epsilon = Double.parseDouble(args[8]);
			slack = Double.parseDouble(args[9]);
			maxStep = Integer.parseInt(args[10]);
			maxProjectionIterations = Integer.parseInt(args[11]);
			System.out.println("Epsilon" + epsilon);
			System.out.println("Slack " + slack);
			System.out.println("Max Step " + maxStep);
			System.out.println("Max projections " + maxProjectionIterations);
		}
	}
	
	 
	 
	 
	
	 
		public void run() throws IOException {
			System.out.println("Loading Corpus");
			forwardCorpus = BilingualCorpus.getCorpusFromFileDescription(corpusDescription, size, maxSentenceLen);
			backwardCorpus = forwardCorpus.reverse();
			modelDir = modelDir + "/" + modelName + "/" + forwardCorpus.getName()
			+ "/" + size + "/model/";
			System.out.println("Finished Loading Corpus");
			System.out.println("Loading Model");
			if (modelName.equalsIgnoreCase("baseline")) {
				mhmm = HMM.loadModel(forwardCorpus, modelDir + "/MHMM/forward");
				mhmmb = HMM.loadModel(backwardCorpus, modelDir + "/MHMM/backward");
			} else if (modelName.equalsIgnoreCase("substochastic")) {
				mhmm = SubstochasticHMM.loadModel(forwardCorpus, modelDir + "/MHMM/forward");
				mhmmb = SubstochasticHMM.loadModel(backwardCorpus, modelDir + "/MHMM/backward");
			}else if (modelName.equalsIgnoreCase("agreement")) {
				mhmm = AgreementHMM.loadModel(forwardCorpus, backwardCorpus,
						modelDir + "/MHMM/");
				mhmmb = ((AgreementHMM) mhmm).getBackwardModel();
			} 
			else if (modelName.equalsIgnoreCase("regularHMM")) {
				mhmm = RegularHMM.loadModel(forwardCorpus, modelDir + "/MHMM/forward");
				mhmmb = RegularHMM.loadModel(backwardCorpus, modelDir + "/MHMM/backward");			
			}else if (modelName.equalsIgnoreCase("symmetric")) {
				mhmm = SymmetricHMM.loadModel(forwardCorpus, backwardCorpus, modelDir + "/MHMM/",epsilon,slack,maxStep,maxProjectionIterations);
				mhmmb = ((SymmetricHMM) mhmm).getBackwardModel(); 
			}else if (modelName.equalsIgnoreCase("bijective")) {
				
				mhmm = BijectiveHMM.loadModel(forwardCorpus, modelDir + "/MHMM/forward",epsilon,slack,maxStep,maxProjectionIterations);
				mhmmb = BijectiveHMM.loadModel(backwardCorpus, modelDir + "/MHMM/backward",epsilon,slack,maxStep,maxProjectionIterations);
				System.out.println("slack " + ((BijectiveHMM) mhmm)._slack);
				System.out.println("epsilon " + ((BijectiveHMM) mhmm)._epsilon);
				System.out.println("max steps " + ((BijectiveHMM) mhmm)._maxStepSize);
				System.out.println("max number iterartions " + ((BijectiveHMM) mhmm)._maxNumberIterations);
			}else {
				System.out.println("Unknown Model: " + modelName);
				System.exit(1);
			}
			System.out.println("Finished Loading Model");
			AlignmentsSet forwardAl = null;
			AlignmentsSet backwardAl = null;
			if (decodingType.equalsIgnoreCase("viter")) {
				forwardAl = mhmm.viterbiAlignments(BilingualCorpus.TEST_CORPUS);
				backwardAl = mhmmb.viterbiAlignments(BilingualCorpus.TEST_CORPUS);
			} else if (decodingType.equalsIgnoreCase("post-aer")) {
				float threshF = mhmm.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,false);
				forwardAl = mhmm.posteriorAlignments(BilingualCorpus.TEST_CORPUS,threshF,false,false);
				float threshB = mhmmb.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,false);
				backwardAl = mhmmb.posteriorAlignments(BilingualCorpus.TEST_CORPUS,threshB,false,false);
			} 
			else if (decodingType.equalsIgnoreCase("post-f1")) {
				float forwardThresh = mhmm.tuneTreshholdF1(
						BilingualCorpus.DEV_CORPUS, posterior,false);
				float backwardThresh = mhmmb.tuneTreshholdF1(
						BilingualCorpus.DEV_CORPUS, posterior,false);
				forwardAl = mhmm.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
						forwardThresh,false,false);
				backwardAl = mhmmb.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
						backwardThresh,false,false);

			} else if (decodingType.startsWith("post-pts")) {
				// sample ~200 sentences from the training corpus
				System.out.println("Forward" + mhmm);
				System.out.println("Backward" + mhmmb);
				Random r = new Random(-9119629771948792409L);
				int totNumSent = forwardCorpus.getNumberOfTrainingSentences();
				int[] sample = new int[200];
				for (int i = 0; i < sample.length; i++) {
					sample[i] = r.nextInt(totNumSent);
				}
				// count number of aligned points with viterbi
				int numAlignmentsF = mhmm.viterbiAlignments(sample,
						BilingualCorpus.TRAIN_CORPUS).numberOfAlignedPoints();
				int numAlignmentsB = mhmmb.viterbiAlignments(sample,
						BilingualCorpus.TRAIN_CORPUS).numberOfAlignedPoints();
				// tune posterior to get that.
				float forwardThresh = mhmm.tuneThresholdNumPoints(numAlignmentsF,
						sample, BilingualCorpus.TRAIN_CORPUS,false);
				float backwardThresh = mhmmb.tuneThresholdNumPoints(numAlignmentsB,
						sample, BilingualCorpus.TRAIN_CORPUS,false);
				forwardAl = mhmm.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
						forwardThresh,false,false);
				backwardAl = mhmmb.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
						backwardThresh,false,false);
			}
			else {
				System.out.println("Unknowon decoding");
				System.exit(-1);
			}
			baseDir = baseDir + "/" + modelName + "/" + forwardCorpus.getName() + "/"
			+ size + "/" + decodingType + "/";
			StaticTools.createDir(baseDir);
			PrintStream forwardPrinter = new PrintStream(baseDir
					+ "/"+projecTestTime+"-forward.tex");
			PrintStream backwardPrinter = new PrintStream(baseDir
					+ "/"+projecTestTime+"backward.tex");
			
			PrintStream bothPrinter = new PrintStream(baseDir
					+ "/"+projecTestTime+"-both.tex");
			
			AlignerOutputLatex.printBegin(forwardPrinter);
			AlignerOutputLatex.outputWithGoldPostLatex(forwardAl, forwardCorpus.getGold(), forwardCorpus, forwardPrinter, 20, "");
			AlignerOutputLatex.printEnd(forwardPrinter);
			
			AlignerOutputLatex.printBegin(backwardPrinter);
			AlignerOutputLatex.outputWithGoldPostLatex(backwardAl, backwardCorpus.getGold(), backwardCorpus, backwardPrinter, 20, "");
			AlignerOutputLatex.printEnd(backwardPrinter);
			
			AlignerOutputLatex.printBegin(bothPrinter);
			
			for(int i = 0; i < forwardCorpus.getGold().size(); i++){
				Alignment forward = forwardAl.get(i);
				Alignment backward = backwardAl.get(i);
				Alignment gold = forwardCorpus.getGold().get(i);
				if(forward._foreignLen < 30 && forward._sourceLen < 30){
					AlignerOutputLatex.outputWithGoldPostLatexForwardAndBackward(forward, backward, gold, forwardCorpus, bothPrinter, "");
				}
			}
			AlignerOutputLatex.printEnd(bothPrinter);
		}
		
	
		

	public static void main(String[] args) throws IOException {
		PrettyPrintAlignments exp = new PrettyPrintAlignments(args);
		exp.run();
	}
	
}
